package com.djf.springaidemo.controller;

import com.alibaba.fastjson.JSON;
import com.djf.springaidemo.memory.MysqlChatMemory;
import com.djf.springaidemo.tools.CsTools;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel; // 修改为正确的包路径
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.bind.DefaultValue;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.UUID;

@RestController
@RequestMapping("/qwen")
public class QwenChatController {

        private final ChatModel chatModel;

        private final MysqlChatMemory chatMemory;

        @Autowired
        @Qualifier("vectorStore2")
        private VectorStore vectorStore;

        public QwenChatController(
                        ChatModel chatModel,
                        MysqlChatMemory mysqlChatMemory) { // 注入 MyBatis-Plus 实现
                this.chatModel = chatModel;
                this.chatMemory = mysqlChatMemory;
        }

        // private final ChatMemory chatMemory = new InMemoryChatMemory();

        String systemPrompt = String.format("""
                        您是虾虾航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                        		您正在通过在线聊天系统与客户互动。
                        		您能够支持已有机票的预订详情查询、机票日期改签、机票预订取消等操作，其余功能将在后续版本中添加，如果用户问的问题不支持请告知详情。
                        	   	在提供有关机票预订详情查询、机票日期改签、机票预订取消等操作之前，您必须始终从用户处获取以下信息：预订号、客户姓名。
                        	   	在询问用户之前，请检查消息历史记录以获取预订号、客户姓名等信息，尽量避免重复询问给用户造成困扰。
                        	   	在更改预订之前，您必须确保条款允许这样做。
                        	   	如果更改需要收费，您必须在继续之前征得用户同意。
                        	   	使用提供的功能获取预订详细信息、更改预订和取消预订。
                        	   	如果需要，您可以调用相应函数辅助完成。
                        	   	请讲中文。
                        	   	今天的日期是 %s.
                        """, new java.util.Date());

        // 同步问答
        @GetMapping("chat")
        public String syncChat(@RequestParam String prompt) {

                return ChatClient.create(chatModel)
                                .prompt()
                                .messages(
                                                new SystemMessage(systemPrompt),
                                                new UserMessage(prompt))
                                .advisors(new MessageChatMemoryAdvisor(chatMemory),
                                                new SimpleLoggerAdvisor(),
                                        new QuestionAnswerAdvisor(vectorStore)
                                )
                                .tools(new CsTools())
                                .call()
                                .content();
        }

        @GetMapping("/stream")
        public Flux<String> streamChat(
                        @RequestParam(value = "prompt", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") String prompt,
                        HttpServletResponse response) {

                response.setCharacterEncoding("UTF-8");
                return ChatClient.create(chatModel)
                                .prompt(prompt)
                                .advisors(new MessageChatMemoryAdvisor(chatMemory))
                                .stream()
                                .content();
        }

        // 流式问答（SSE）
        @GetMapping(value = "chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
        public Flux<ServerSentEvent<String>> streamChat(@RequestParam String prompt) {
                return ChatClient.create(chatModel)
                                .prompt()
                                .messages(
                                                new SystemMessage("你是一个Spring专家，帮助用户回答Java/Spring框架问题"),
                                                new UserMessage(prompt))
                                .stream()
                                .chatResponse()
                                .map(chatResponse -> {
                                        return ServerSentEvent.builder(
                                                        JSON.toJSONString(chatResponse)).event("message").build();
                                });
        }

        @GetMapping(value = "chat/stream/memory", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
        public Flux<ServerSentEvent<String>> memoryChat(
                        @RequestParam String prompt,
                        @RequestParam(required = false) String sessionId) {

                if (sessionId == null || sessionId.isBlank()) {
                        sessionId = UUID.randomUUID().toString();
                }

                var advisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 5);

                return ChatClient.create(chatModel)
                                .prompt()
                                .user(prompt)
                                .advisors(advisor)
                                .stream()
                                .chatResponse()
                                .map(chatResponse -> ServerSentEvent.<String>builder()
                                                .data(JSON.toJSONString(chatResponse))
                                                .event("message")
                                                .build());
        }

        /*
         * // 带记忆的流式问答
         * 
         * @GetMapping(value = "chat/stream/memory", produces =
         * MediaType.APPLICATION_JSON_UTF8_VALUE)
         * public Flux<ServerSentEvent<String>> memoryChat(
         * 
         * @RequestParam String prompt,
         * 
         * @RequestParam(required = false) String sessionId) {
         * 
         * // 如果未提供 sessionId，则生成一个随机的 conversationId
         * if (sessionId == null || sessionId.isBlank()) {
         * sessionId = UUID.randomUUID().toString();
         * }
         * 
         * var advisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 5);
         * 
         * return ChatClient.create(chatModel)
         * .prompt()
         * .user(prompt)
         * .advisors(advisor)
         * .stream()
         * .chatResponse()
         * .map(chatResponse -> ServerSentEvent.builder(
         * JSON.toJSONString(chatResponse))
         * .event("message")
         * .build());
         * }
         */

}
